/*
 * auth.c
 *
 *  Created on: Feb 12, 2012
 *      Author: jiabin
 */

#include "auth.h"
#include "splitter.h"
#include "secure_rw.h"
#include <sys/types.h>
#include <sys/stat.h>
#include <fcntl.h>
#include <unistd.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <openssl/bio.h>
#include <openssl/hmac.h>
#include <openssl/aes.h>
#include <openssl/x509.h>
#include <openssl/pem.h>

static const char *ca_cert = "ca_cert.pem";

#define BUF_SIZE 2048

static void error(const char *s)
{
	perror(s);
	exit(1);
}

static void get_subject_CN(X509 *cert, unsigned char **buf)
{
	X509_NAME *name = X509_get_subject_name(cert);
	int index = X509_NAME_get_index_by_NID(name, NID_commonName, -1);
	X509_NAME_ENTRY *entry = X509_NAME_get_entry(name, index);

	ASN1_STRING *str = X509_NAME_ENTRY_get_data(entry);

	ASN1_STRING_to_UTF8(buf, str);
}

static enum bool certificate_check(X509 *cert, const char *tarCN, X509 *ca_cert)
{
	/* verify common name */
	unsigned char *cn;
	get_subject_CN(cert, &cn);

	if (strcmp((char *) cn, tarCN) != 0)
	{
		fprintf(stderr, "Wrong Common Name!\n");
		return FALSE;
	}

	EVP_PKEY *ca_pub_key = X509_get_pubkey(ca_cert);

	OpenSSL_add_all_digests();
	int ret = X509_verify(cert, ca_pub_key);
	if (ret != 1)
	{
		fprintf(stderr, "Invalid signature!\n");
		return FALSE;
	}

	fprintf(stderr, "%s's cert verified!\n", cn);
	return TRUE;
}

static void send_cert(const char *cert_file, int sockfd)
{
	unsigned char buf[BUF_SIZE];
	bzero(buf, BUF_SIZE);

	int fd = open(cert_file, O_RDONLY);
	if (fd < 0)
		error("Can't open own_cert");

	int len = read(fd, buf, BUF_SIZE);
	if (len < 0)
		error("Can't read own_cert");
	close(fd);

	write(sockfd, buf, len);
}

static X509 *recv_cert(int sockfd)
{
	unsigned char buf[BUF_SIZE];
	bzero(buf, BUF_SIZE);

	int len = read(sockfd, buf, BUF_SIZE);
	BIO *mem = BIO_new_mem_buf(buf, len);
	X509 *tar_cert = PEM_read_bio_X509(mem, NULL, NULL, NULL);
	return tar_cert;
}

static unsigned char *rand_128_bit()
{
	unsigned char *ret = malloc(16);
	int i;
	for (i = 0; i < 8; ++i)
	{
		int r = rand();
		memcpy(ret + 2 * i, &r, 2);
	}

	return ret;
}

static RSA *get_own_rsa(const char *own_key_file)
{
	FILE *fp = fopen(own_key_file, "r");
	EVP_PKEY *key = PEM_read_PrivateKey(fp, NULL, NULL, NULL);
	RSA *rsa = EVP_PKEY_get1_RSA(key);
	fclose(fp);
	return rsa;
}

static RSA *get_tar_rsa(X509 *tar_cert)
{
	EVP_PKEY *key = X509_get_pubkey(tar_cert);
	RSA *rsa = EVP_PKEY_get1_RSA(key);
	return rsa;
}

static int certificate_clone(X509 *ori_cert, EVP_PKEY *pkey, const char *ca_key_file,
		const char *output_file)
{
	int ret = X509_set_pubkey(ori_cert, pkey);
	if (ret == 0)
	{
		fprintf(stderr, "Replace public key failed!\n");
		return -1;
	}

	FILE *fp_ca_key = fopen(ca_key_file, "r");
	EVP_PKEY *ca_private_key = PEM_read_PrivateKey(fp_ca_key, NULL, NULL, NULL);
	fclose(fp_ca_key);

	X509_sign(ori_cert, ca_private_key, EVP_md5());

	FILE *fp_output = fopen(output_file, "w");
	PEM_write_X509(fp_output, ori_cert);
	fclose(fp_output);

	return 0;
}

static int encrypt_and_send(int fd, RSA *pub_key, void *content, size_t length)
{
	unsigned char buf[256];
	int n = RSA_public_encrypt(length, content, buf, pub_key,
			RSA_PKCS1_PADDING);
	if (n < 256)
	{
		fprintf(stderr, "Error during encrypt and send!\n");
		return -1;
	}
	n = write(fd, buf, 256);
	if (n < 256)
	{
		fprintf(stderr, "Error during encrypt and send!\n");
		return -1;
	}

	return n;
}

static int read_and_decrypt(int fd, RSA *private_key, void *content,
		size_t length)
{
	unsigned char buf[256];
	int n = read(fd, buf, 256);
	if (n < 256)
	{
		fprintf(stderr, "Error during read_and_decrypt!\n");
		return -1;
	}
	n = RSA_private_decrypt(256, buf, content, private_key, RSA_PKCS1_PADDING);
	if (n < length)
	{
		fprintf(stderr, "Error during read_and_decrypt!\n");
		return -1;
	}
	return n;
}

int auth_active(struct connection *conn, const char *own_cert_file,
		const char *own_key_file, const char *tarCN)
{
	FILE *fp_ca = fopen(ca_cert, "r");
	X509 *ca_cert = PEM_read_X509(fp_ca, NULL, NULL, NULL);
	fclose(fp_ca);

	/* 1. Alice sends her own certificate to Bob. */
	send_cert(own_cert_file, conn->sockfd);

	/* 4. Alice verifies Bob's certificate by check CN and signature. */
	X509 *tar_cert = recv_cert(conn->sockfd);
	if (!certificate_check(tar_cert, tarCN, ca_cert))
		return -1;

	RSA *own_rsa = get_own_rsa(own_key_file);
	RSA *tar_rsa = get_tar_rsa(tar_cert);

	/* 5. Alice generates a random integer x, encrypts it using Bob's public key, sends it to Bob. */
	int r = rand();
	if (encrypt_and_send(conn->sockfd, tar_rsa, &r, sizeof(int)) < 0)
		return -1;

	/* 7. Alice decrypts and verifies x - 1. */
	int r_recv;
	if (read_and_decrypt(conn->sockfd, own_rsa, &r_recv, sizeof(int)) < 0)
		return -1;
	if (r_recv != r - 1)
		return -1;

	/* 9. Alice decrypts the number x, substract it by one, encrypts using Bob's public key, send it back to Bob. */
	if (read_and_decrypt(conn->sockfd, own_rsa, &r, sizeof(int)) < 0)
		return -1;
	r--;
	if (encrypt_and_send(conn->sockfd, tar_rsa, &r, sizeof(int)) < 0)
		return -1;

	/* 11. Alice generates a random 128-bit string as AES key, encrypts it using Bob's public key, sends it to Bob. */
	unsigned char *secret = rand_128_bit();
	if (encrypt_and_send(conn->sockfd, tar_rsa, secret, SECRET_BITS / 8) < 0)
		return -1;

	/* 14. Alice decryptes the IV using her own private key. */
	unsigned char *iv = malloc(16);
	if (read_and_decrypt(conn->sockfd, own_rsa, iv, SECRET_BITS / 8) < 0)
		return -1;

	conn->secret = secret;
	conn->iv = iv;

	fprintf(stderr, "Key exchanged!\n");

	return 0;
}

int auth_passive(struct connection *conn, const char *own_cert_file,
		const char *own_key_file, const char *tarCN, enum bool forge)
{
	FILE *fp_ca = fopen(ca_cert, "r");
	X509 *ca_cert = PEM_read_X509(fp_ca, NULL, NULL, NULL);
	fclose(fp_ca);

	/* 2. Bob verifies Alice's certificate by check CN and signature. */
	X509 *tar_cert = recv_cert(conn->sockfd);
	if (!certificate_check(tar_cert, tarCN, ca_cert))
		return -1;

	/* 3. Bob sends his own certificate to Alice. */
	send_cert(own_cert_file, conn->sockfd);

	RSA *own_rsa = get_own_rsa(own_key_file);
	RSA *tar_rsa = get_tar_rsa(tar_cert);

	/* 6. Bob decrypts the number x, substract it by one, encrypts using Alice's public key, send it back to Alice. */
	int r;
	if (read_and_decrypt(conn->sockfd, own_rsa, &r, sizeof(int)) < 0)
		return -1;
	r--;
	if (encrypt_and_send(conn->sockfd, tar_rsa, &r, sizeof(int)) < 0)
		return -1;

	/* 8. Bob generates a random integer x, encrypts it using Alice's public key, sends it to Alice. */
	r = rand();
	if (encrypt_and_send(conn->sockfd, tar_rsa, &r, sizeof(int)) < 0)
		return -1;

	/* 10. Bob decrypts and verifies x - 1. */
	int r_recv;
	if (read_and_decrypt(conn->sockfd, own_rsa, &r_recv, sizeof(int)) < 0)
		return -1;
	if (r_recv != r - 1)
		return -1;

	/* 12. Bob decrypts the AES key using his own private key. */
	unsigned char *secret = malloc(16);
	if (read_and_decrypt(conn->sockfd, own_rsa, secret, SECRET_BITS / 8) < 0)
		return -1;

	/* 13. Bob generates a random 128-bit string as IV, encrypts it using Alice's public key, sends it to Alice. */
	unsigned char *iv = rand_128_bit();
	if (encrypt_and_send(conn->sockfd, tar_rsa, iv, SECRET_BITS / 8) < 0)
		return -1;

	conn->secret = secret;
	conn->iv = iv;

	fprintf(stderr, "Key exchanged!\n");

	/* 15. If Bob wants to forge Alice's certificate, he does so by replacing it with his own public key,
	 *     assuming he has access to CA's private key to sign the modified certificate. */
	if (forge)
	{
		FILE *fp_own_cert = fopen(own_cert_file, "r");
		X509 *own_cert = PEM_read_X509(fp_own_cert, NULL, NULL, NULL);
		EVP_PKEY *own_key = X509_get_pubkey(own_cert);
		int r = certificate_clone(tar_cert, own_key, "ca_key.pem", "forge.pem");
		if (r < 0)
			return -1;
	}

	return 0;
}
